/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: alg remote update.
 * Author: qingjicheng
 * Create: 2022-11-28
 */

#include "alg_coll_remote_update.h"
#include "log.h"
#include "dtype_common.h"
#include "adapter_hccp.h"
#include "mr_manager.h"
#include "profiling_manager.h"
#include "adapter_prof.h"
#include "sal_pub.h"

using namespace std;
constexpr u32 SMALL_PAGE_SIZE = 4096;
namespace hccl {
AlgCollRemoteUpdate::AlgCollRemoteUpdate() : funcTable_({
    {EmbeddingServiceParam::OperatorType::COLL_REMOTE_UPDATE,
        std::bind(&AlgCollRemoteUpdate::AlgRemoteUpdate, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::COLL_REMOTE_UPDATE_PAIRED,
        std::bind(&AlgCollRemoteUpdate::AlgRemoteUpdate, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::RECV_UPDATE_REQUEST,
        std::bind(&AlgCollRemoteUpdate::AlgRecvUpdateRequest, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::SEND_UPDATE_RESPONSE,
        std::bind(&AlgCollRemoteUpdate::AlgIsendUpdateResponse, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::WAIT_SOME,
        std::bind(&AlgCollRemoteUpdate::AlgWaitSome, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::SERVICE_CANCEL,
        std::bind(&AlgCollRemoteUpdate::AlgServiceCancel, this, std::placeholders::_1)}
})
{
}

AlgCollRemoteUpdate::~AlgCollRemoteUpdate()
{
}

HcclResult AlgCollRemoteUpdate::Process(HcclOpPtr opPtr)
{
    CHK_PTR_NULL(opPtr);
    auto op = static_cast<HcclOperator<EmbeddingServiceParam> *>(opPtr);
    CHK_PRT_RET(funcTable_.find(op->param_.opType) == funcTable_.end(),
        HCCL_ERROR("[Process]funcTable_ without opType[%u]", op->param_.opType), HCCL_E_PARA);
    CHK_RET(funcTable_[op->param_.opType](op));
    return HCCL_SUCCESS;
}

HcclResult AlgCollRemoteUpdate::AlgRemoteUpdate(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    return HCCL_SUCCESS;
}

HcclResult AlgCollRemoteUpdate::AllocTransferMemAndRegMrForPs(const HcclOperator<EmbeddingServiceParam> &op)
{
    if (transferMemAllocated_) {
        HCCL_DEBUG("key and Value TransferMem has been Allocated.");
        return HCCL_SUCCESS;
    }

    if (keyTransferBeginMem_ == nullptr) {
        u64 len = op.param_.keyMaxNum * SIZE_TABLE[op.param_.keyType];
        // 申请内存对齐
        u64 memLen = len + SMALL_PAGE_SIZE;

        EXECEPTION_CATCH((keyTransferBeginMem_ = make_unique<s8[]>(memLen)), return HCCL_E_PTR);
        CHK_PTR_NULL(keyTransferBeginMem_);

        keyTransferMem_ = reinterpret_cast<void *>(Align<u64>(reinterpret_cast<u64>(keyTransferBeginMem_.get()),
            SMALL_PAGE_SIZE));

        HCCL_INFO("keyTransferBeginMem[%llu], memLen[%llu], keyTransferMem[%llu], len[%llu] GlobalMr Registered",
            hash<s8 *>{}(keyTransferBeginMem_.get()), memLen, hash<void *>{}(keyTransferMem_), len);
    }

    if (valueTransferBeginMem_ == nullptr) {
        u64 len = op.param_.valueCount * SIZE_TABLE[op.param_.valuesType];
        // 申请内存对齐
        u64 memLen = len + SMALL_PAGE_SIZE;

        EXECEPTION_CATCH((valueTransferBeginMem_ = make_unique<s8[]>(memLen)), return HCCL_E_PTR);
        CHK_PTR_NULL(valueTransferBeginMem_);

        valueTransferMem_ = reinterpret_cast<void *>(Align<u64>(reinterpret_cast<u64>(valueTransferBeginMem_.get()),
            SMALL_PAGE_SIZE));

        HCCL_INFO("valueTransferBeginMem[%llu], memLen[%llu], valueTransferMem[%llu], len[%llu] GlobalMr Registered",
            hash<s8 *>{}(valueTransferBeginMem_.get()), memLen, hash<void *>{}(valueTransferMem_), len);
    }

    transferMemAllocated_ = true;

    return HCCL_SUCCESS;
}

HcclResult AlgCollRemoteUpdate::AlgRecvUpdateRequest(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    uint64_t beginTime = hrtMsprofSysCycleTime();

    auto maxValue = *max_element(opPtr->param_.getRequestParam.workerList.begin(),
        opPtr->param_.getRequestParam.workerList.end());
    opPtr->param_.embeddingParam.keysAddrInfo.resize(maxValue + 1);
    opPtr->param_.embeddingParam.valuesAddrInfo.resize(maxValue + 1);
    opPtr->param_.getRequestParam.rankTransKeysCompFlag.resize(maxValue + 1, -1);
    opPtr->param_.getRequestParam.rankTransValueCompFlag.resize(maxValue + 1, -1);
    for (auto iter : opPtr->param_.getRequestParam.workerList) {
        opPtr->param_.getRequestParam.rankTransKeysCompFlag[iter] = 1;
        opPtr->param_.getRequestParam.rankTransValueCompFlag[iter] = 1;
    }

    {
        // KEY PROBE
        UpdateTransTask *transTaskPtr = new (std::nothrow) UpdateTransTask();
        CHK_PTR_NULL(transTaskPtr);
        transTaskPtr->improbeRanks = opPtr->param_.getRequestParam.workerList;
        transTaskPtr->handle = &opPtr->param_.getRequestParam.updateKeysSize;
        transTaskPtr->addrInfo = &opPtr->param_.embeddingParam.keysAddrInfo;
        transTaskPtr->rankTransCompPtr = &opPtr->param_.getRequestParam.rankTransKeysCompFlag;
        transTaskPtr->outputs.push_back(opPtr->param_.getRequestParam.reqStatus);

        // rdma情况下有效
        transTaskPtr->rdmaResponseAddrsPtr = &opPtr->param_.getRequestParam.rdmaResponseAddrs;

        transTaskPtr->taskType = OpTaskType::PROBE;
        for (auto &rank : transTaskPtr->improbeRanks) {
            transTaskPtr->recvWaitFlagMap[rank] = true;
        }
        OpTaskPtr taskPtr(transTaskPtr);
        opPtr->PushTask(std::move(taskPtr));
    }

    {
        // VALUE PROBE
        UpdateTransTask *transTaskPtr = new (std::nothrow) UpdateTransTask();
        CHK_PTR_NULL(transTaskPtr);
        transTaskPtr->improbeRanks = opPtr->param_.getRequestParam.workerList;
        transTaskPtr->handle = &opPtr->param_.getRequestParam.updateValuesSize;
        transTaskPtr->addrInfo = &opPtr->param_.embeddingParam.valuesAddrInfo;
        transTaskPtr->rankTransCompPtr = &opPtr->param_.getRequestParam.rankTransValueCompFlag;

        // rdma情况下有效
        transTaskPtr->rdmaResponseAddrsPtr = &opPtr->param_.getRequestParam.rdmaResponseAddrs;

        transTaskPtr->taskType = OpTaskType::PROBE;
        for (auto &rank : transTaskPtr->improbeRanks) {
            transTaskPtr->recvWaitFlagMap[rank] = true;
        }
        OpTaskPtr taskPtr(transTaskPtr);
        opPtr->PushTask(std::move(taskPtr));
    }

    opPtr->param_.embeddingParam.keyTransferMem = keyTransferMem_;
    opPtr->param_.embeddingParam.valueTransferMem = valueTransferMem_;

    HCCL_DEBUG("keyTransferMem[%llu], valueTransferMem[%llu]", hash<void *>{}(keyTransferMem_),
        hash<void *>{}(valueTransferMem_));

    // key & value IRECV
    u32 pNum = opPtr->param_.getRequestParam.psMap.size();
    u32 wNum = opPtr->param_.getRequestParam.workerList.size();
    for (u32 i = 0; i < pNum; i++) {
        for (auto &iter : opPtr->param_.getRequestParam.psMap) {
            if (opPtr->param_.getRequestParam.rankId != iter.second) {
                continue;
            }
            u32 dstWorkIdxTmp = (iter.first + i) % pNum;
            u32 j = 0;
            while (dstWorkIdxTmp + j < wNum) {
                u32 dstWorkIdx = (dstWorkIdxTmp + j) % wNum;
                {
                    // KEY IRECV
                    UpdateTransTask *transTaskPtr = new (std::nothrow) UpdateTransTask();
                    CHK_PTR_NULL(transTaskPtr);
                    transTaskPtr->dstRank = opPtr->param_.getRequestParam.workerList.at(dstWorkIdx);
                    transTaskPtr->recvMem = static_cast<u8 *>(opPtr->param_.embeddingParam.keyTransferMem) +
                        opPtr->param_.keyMaxNum / wNum * SIZE_TABLE[opPtr->param_.keyType] * dstWorkIdx;
                    transTaskPtr->transferMemBegin = static_cast<u8 *>(opPtr->param_.embeddingParam.keyTransferMem);
                    transTaskPtr->transferMemLen = opPtr->param_.keyMaxNum * SIZE_TABLE[opPtr->param_.keyType];
                    transTaskPtr->dataType = opPtr->param_.keyType;
                    transTaskPtr->count = opPtr->param_.keyMaxNum / wNum;
                    transTaskPtr->taskType = OpTaskType::IRECV;
                    transTaskPtr->needRecordFlag = false;
                    OpTaskPtr taskPtr(transTaskPtr);
                    opPtr->param_.embeddingParam.keysAddrInfo[transTaskPtr->dstRank].startAddr = transTaskPtr->recvMem;
                    HCCL_INFO("[Alg][RecvUpdateRequest]: rank[%u] keyRecvMem[%p]",
                        transTaskPtr->dstRank, transTaskPtr->recvMem);
                    opPtr->PushTask(std::move(taskPtr));
                }

                {
                    // VALUE IRECV
                    UpdateTransTask *transTaskPtr = new (std::nothrow) UpdateTransTask();
                    CHK_PTR_NULL(transTaskPtr);
                    transTaskPtr->dstRank = opPtr->param_.getRequestParam.workerList.at(dstWorkIdx);
                    transTaskPtr->recvMem = static_cast<u8 *>(opPtr->param_.embeddingParam.valueTransferMem) +
                        opPtr->param_.valueCount / wNum * SIZE_TABLE[opPtr->param_.valuesType] * dstWorkIdx;
                    transTaskPtr->transferMemBegin = static_cast<u8 *>(opPtr->param_.embeddingParam.valueTransferMem);
                    transTaskPtr->transferMemLen = opPtr->param_.valueCount * SIZE_TABLE[opPtr->param_.valuesType];
                    transTaskPtr->dataType = opPtr->param_.valuesType;
                    transTaskPtr->count = opPtr->param_.valueCount / wNum;
                    transTaskPtr->taskType = OpTaskType::IRECV;
                    // 目前无流水等，worker不再接收这个notify了，故改为false
                    transTaskPtr->needRecordFlag = false;
                    OpTaskPtr taskPtr(transTaskPtr);
                    opPtr->param_.embeddingParam.valuesAddrInfo[transTaskPtr->dstRank].startAddr =
                        transTaskPtr->recvMem;
                    HCCL_INFO("[Alg][RecvUpdateRequest]: rank[%u] valueRecvMem[%p]", transTaskPtr->dstRank,
                        transTaskPtr->recvMem);
                    opPtr->PushTask(std::move(taskPtr));
                }

                j += pNum;
            }
        }
    }

    {
        // WAIT
        WaitSomeTask *transTaskPtr = new (std::nothrow) WaitSomeTask();
        CHK_PTR_NULL(transTaskPtr);
        transTaskPtr->taskType = OpTaskType::WAIT;
        OpTaskPtr taskPtr(transTaskPtr);
        opPtr->PushTask(std::move(taskPtr));
    }

    // 全局key-value update
    TIME_PRINT(CHK_RET(AlgUpdate(opPtr)));

    auto &profilingManager = hccl::ProfilingManager::Instance();
    CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime, ProfTaskType::TASK_UPDATE_ALG_GLOBAL_REDUCE));
    return HCCL_SUCCESS;
}

HcclResult AlgCollRemoteUpdate::AlgUpdate(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    PsReduceSumTask *reduceSumTaskkPtr = new (std::nothrow) PsReduceSumTask();
    CHK_PTR_NULL(reduceSumTaskkPtr);
    reduceSumTaskkPtr->keyMaxNum = opPtr->param_.keyMaxNum;
    reduceSumTaskkPtr->keys = static_cast<s64 *>(opPtr->param_.keys);
    reduceSumTaskkPtr->value = opPtr->param_.values;
    reduceSumTaskkPtr->taskType = OpTaskType::REDUCE_SUM;
    reduceSumTaskkPtr->keyTransferMem = opPtr->param_.embeddingParam.keyTransferMem;
    reduceSumTaskkPtr->valueTransferMem = opPtr->param_.embeddingParam.valueTransferMem;

    reduceSumTaskkPtr->rdmaEnveInfosTransferMem = opPtr->param_.embeddingParam.rdmaEnveInfosTransferMem;
    reduceSumTaskkPtr->haveRdmaConn = opPtr->param_.haveRdmaConn;

    reduceSumTaskkPtr->updateKeysSize = &opPtr->param_.getRequestParam.updateKeysSize;
    reduceSumTaskkPtr->updateValuesSize = &opPtr->param_.getRequestParam.updateValuesSize;
    reduceSumTaskkPtr->actualKeyCount = &(opPtr->param_.getRequestParam.updateReqStatus->actualKeyCount);
    reduceSumTaskkPtr->actualValueCount = &(opPtr->param_.getRequestParam.updateReqStatus->actualValueCount);
    reduceSumTaskkPtr->keyType = opPtr->param_.keyType;
    reduceSumTaskkPtr->valueType = opPtr->param_.valuesType;
    reduceSumTaskkPtr->keyAddrInfo = &opPtr->param_.embeddingParam.keysAddrInfo;
    reduceSumTaskkPtr->valueAddrInfo = &opPtr->param_.embeddingParam.valuesAddrInfo;
    reduceSumTaskkPtr->workNum = opPtr->param_.getRequestParam.workerList.size();

    u64 realValueCount = opPtr->param_.valueCount * SIZE_TABLE[opPtr->param_.valuesType];

    HCCL_INFO("keyMaxNum[%u] realValueBytes[%llu]", opPtr->param_.keyMaxNum, realValueCount);
    OpTaskPtr taskPtr(reduceSumTaskkPtr);
    opPtr->PushTask(std::move(taskPtr));
    return HCCL_SUCCESS;
}

HcclResult AlgCollRemoteUpdate::AlgIsendUpdateResponse(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    *opPtr->param_.getRequestParam.outRequest = nullptr;
    CHK_PRT(SendNoSignalTask(opPtr->param_.getRequestParam.workerList, opPtr->param_.getRequestParam.transferMems,
        opPtr, OpTaskType::ISEND, keyTransferMem_));

    return HCCL_SUCCESS;
}

HcclResult AlgCollRemoteUpdate::AlgWaitSome(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    WaitSomeTask *transTaskPtr = new (std::nothrow) WaitSomeTask();
    CHK_PTR_NULL(transTaskPtr);
    transTaskPtr->requestCount = opPtr->param_.waitSomeParam.requestCount;
    transTaskPtr->requestArray = opPtr->param_.waitSomeParam.requestArray;
    transTaskPtr->compCount = opPtr->param_.waitSomeParam.compCount;
    transTaskPtr->compIndices = opPtr->param_.waitSomeParam.compIndices;
    transTaskPtr->compStatus = opPtr->param_.waitSomeParam.compStatus;
    transTaskPtr->taskType = OpTaskType::WAIT;
    OpTaskPtr taskPtr(transTaskPtr);
    opPtr->PushTask(std::move(taskPtr));
    return HCCL_SUCCESS;
}
} // namespace hccl
